"""
预标签器

使用大模型对PR修改数据进行预标签，为后续的few-shot选择提供标签信息
"""

import json
import logging
import os
import time
from typing import Dict, Any, List, Optional, Tuple
from openai import OpenAI
from prompts import create_prelabel_prompt
from config import Config

# 获取日志记录器
logger = logging.getLogger(__name__)

# 文档编辑标签定义
DOCUMENT_EDIT_LABELS = [
    # ===== 拼写类 =====
    {"label_id": 1, "type": "拼写-英文专业名词大小写", "detail": "英文专业术语、单位、品牌名等的大小写不规范，需要统一为标准形式，如cpu→CPU、api→API、github→GitHub、8gb→8GB、Install.md→install.md", "unit": "句子"},
    {"label_id": 2, "type": "拼写-英文同词大小写一致", "detail": "同一章节内相同的非专业英文单词大小写不一致，需要统一为同一形式，如test/Test/TEST在同一章节内应统一为test", "unit": "章节"},
    {"label_id": 3, "type": "拼写-中文相似字错误", "detail": "中文词汇使用了同音字或形近字，需要替换为正确汉字，如区块连→区块链、阀值→阈值、布署→部署", "unit": "句子"},
    {"label_id": 4, "type": "拼写-中文误触多余字", "detail": "因键盘误操作产生的重复字符，需要删除多余字符，如安装包包→安装包、打开开文档→打开文档", "unit": "句子"},

    # ===== 超链接类 =====
    {"label_id": 5, "type": "超链接-描述与格式错误", "detail": "超链接的文字描述与目标不符、格式不规范或地址错误，需要更正链接描述、补充格式标记或修正地址", "unit": "超链接"},

    # ===== 空格类 =====
    {"label_id": 6, "type": "空格-中英文与标点空格错误", "detail": "空格使用不规范，包括纯中文间的多余空格、中英文混排时英文前后缺少空格、中英文标点前后空格不当等问题", "unit": "句子/代码行"},

    # ===== 标点类 =====
    {"label_id": 7, "type": "标点-中英文标点使用错误", "detail": "中英文标点符号混用，如英文代码中使用中文标点、中文文本中使用英文标点，需要统一标点类型", "unit": "代码行/句子"},
    {"label_id": 8, "type": "标点-顿号使用错误", "detail": "顿号使用场景不当，应该仅在句子内并列词语间使用顿号，其他场景需改用适当标点", "unit": "句子"},
    {"label_id": 9, "type": "标点-成对标点缺失", "detail": "成对标点符号缺少一侧，需要补齐缺失的引号、括号、书名号等", "unit": "句子/段落"},
    {"label_id": 10, "type": "标点-同级内容标点一致", "detail": "同级内容的标点使用不统一，需要调整为相同的标点规则，确保段落结尾、列表项等标点保持一致", "unit": "段落/表格/列表"},

    # ===== 格式 & 版式类 =====
    {"label_id": 11, "type": "格式-换行不一致", "detail": "同级内容的换行格式不统一，需要调整为相同的换行逻辑和缩进规范", "unit": "段落"},
    {"label_id": 12, "type": "版式-视觉排版优化", "detail": "文档中存在影响视觉一致性和阅读体验的排版问题，包括图片位置不当（如遮挡文字、偏离语义位置）、行间距/段落间距不统一、对齐方式混乱等，需调整布局以提升可读性", "unit": "段落/文档"},
    # ===== 标记与代码类 =====
    {"label_id": 13, "type": "标记-文档标记规范化", "detail": "文档标记语法错误或格式不统一，包括Markdown、HTML、XML标签语法错误，格式风格不一致，文档内路径描述错误等", "unit": "文档/代码块"},
    {"label_id": 14, "type": "标记-代码块规范", "detail": "代码块缺语言标识、命令语法错误、注释符号错误等（如“```→```python”“//注释→#注释”）", "unit": "代码块"},

    # ===== 文件与路径类 =====
    {"label_id": 15, "type": "文件-命名规范", "detail": "文件名语义不当或扩展名错误，需要更正文件名含义或修正扩展名类型，不含文件名大小写错误", "unit": "文件"},
    {"label_id": 16, "type": "文件-路径规范", "detail": "实际文件路径错误，需要更正为正确的文件系统路径", "unit": "文件"},

    # ===== 表达类 =====
    {"label_id": 17, "type": "表达-语义优化", "detail": "句子表达不够准确或流畅，需要优化语义表达以提升内容的准确性和可读性", "unit": "句子/段落"},
    {"label_id": 18, "type": "表达-冗余精简", "detail": "存在语义重复或冗余表达，需要删除重复内容并精简表述", "unit": "句子/段落"},

    # ===== 结构类 =====
    {"label_id": 19, "type": "结构-章节顺序优化", "detail": "章节编号或逻辑顺序不合理，需要调整为正确的章节序号和逻辑结构", "unit": "章节/文档"},

    # ===== 标题类 =====
    {"label_id": 20, "type": "标题-语义优化", "detail": "标题语义不够准确或描述性不强，需要优化标题内容使其更好地概括章节内容", "unit": "标题"},

    # ===== 信息类=====
    {"label_id": 21, "type": "信息-完整性补充", "detail": "信息不完整或缺失关键内容，需要补充操作步骤、参数说明、技术背景等必要信息", "unit": "句子/段落"},

    {"label_id": 22, "type": "其他", "detail": "不属于以上任何类别的情况", "unit": "文档/句子"}
]


class PreLabeler:
    """预标签器"""
    
    def __init__(self, api_key: str, base_url: str, model_name: str = None):
        """
        初始化预标签器
        
        Args:
            api_key: API密钥
            base_url: API基础URL
            model_name: 模型名称
        """
        self.api_key = api_key
        # 确保base_url格式正确，拼接完整的API端点
        self.base_url = base_url.rstrip("/") + "/v1/chat/completions"
        self.model_name = model_name or Config.get_model_name()
        
        # 初始化OpenAI客户端
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url.rstrip("/") + "/v1"  # OpenAI客户端只需要到/v1
        )
        
        self.allowed_label_ids = {
            tag['label_id']
            for tag in DOCUMENT_EDIT_LABELS
            if isinstance(tag.get('label_id'), int)
        }
        self.label_id_to_name = {
            tag['label_id']: (tag.get('type') or f"label_{tag['label_id']}")
            for tag in DOCUMENT_EDIT_LABELS
            if isinstance(tag.get('label_id'), int)
        }

        # 统计信息
        self.total_processed = 0
        self.total_filtered = 0
        self.label_distribution = {}
        
        logger.info(f"初始化预标签器，模型: {model_name}")
    
    def create_labeling_prompt(self, data: Dict[str, Any]) -> str:
        """创建标签提示词（统一从 prompts 构造）"""
        return create_prelabel_prompt(data, DOCUMENT_EDIT_LABELS)
    
    def call_model_for_labeling(self, prompt: str) -> Optional[Dict[str, Any]]:
        """调用模型进行标签"""
        import time as _t
        start_ts = _t.time()
        prompt_len = len(prompt or '')
        try:
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=2048
            )
            
            content = response.choices[0].message.content.strip()
            duration = _t.time() - start_ts
            
            # 尝试解析JSON响应
            try:
                # 处理可能的Markdown格式（包含```json和```标记）
                content_clean = content.strip()
                if content_clean.startswith('```json'):
                    # 移除开头的```json
                    content_clean = content_clean[7:]
                if content_clean.startswith('```'):
                    # 移除开头的```
                    content_clean = content_clean[3:]
                if content_clean.endswith('```'):
                    # 移除结尾的```
                    content_clean = content_clean[:-3]
                
                # 清理后的内容
                content_clean = content_clean.strip()
                
                result = json.loads(content_clean)
                return result
            except json.JSONDecodeError:
                # 如果不是JSON格式，检查是否表示不属于任何标签
                if "不属于任何标签类型" in content:
                    # 返回特殊标记，包含模型的分类原因
                    return {"skip_reason": content, "raw_response": content}
                else:
                    logger.warning(f"模型返回格式异常: {content}")
                    logger.debug(f"清理后的内容: {content_clean}")
                    return None
                    
        except Exception as e:
            duration = _t.time() - start_ts
            logger.error(f"调用模型失败: {e}")
            return None
    
    def process_single_data(self, data: Dict[str, Any]) -> Tuple[bool, Optional[Dict[str, Any]], Optional[str]]:
        """处理单条数据
        
        Returns:
            (is_valid, processed_data, skip_reason): 
            - is_valid: 是否有效（True=有效标签，False=跳过或无效）
            - processed_data: 处理后的数据（如果有效）
            - skip_reason: 跳过的原因（如果跳过）
        """
        self.total_processed += 1
        
        # 创建提示词
        prompt = self.create_labeling_prompt(data)
        
        # 调用模型
        label_result = self.call_model_for_labeling(prompt)
        
        if label_result and 'label_id' in label_result:
            # 验证标签ID是否在允许范围内
            label_id = label_result['label_id']
            # 兼容字符串数字
            if isinstance(label_id, str) and label_id.isdigit():
                label_id = int(label_id)

            if label_id in self.allowed_label_ids:
                # 更新统计信息
                label_name = label_result.get('label_name') or self.label_id_to_name.get(label_id, f'label_{label_id}')
                # 若模型未返回label_name则补全
                label_result['label_name'] = label_name
                self.label_distribution[label_name] = self.label_distribution.get(label_name, 0) + 1
                
                # 添加标签信息到原始数据
                data['pre_label'] = label_result

                # 记录成功打标日志
                try:
                    logger.info(
                        f"预标签成功: unit_id={data.get('unit_id','N/A')}, pr={data.get('pr_number','N/A')}, "
                        f"file={data.get('file_path','N/A')} -> label_id={label_id}, label_name={label_name}, "
                        f"reason={label_result.get('reason','')}"
                    )
                except Exception:
                    pass
                return True, data, None
            else:
                logger.warning(
                    f"标签ID {label_id} 不在允许范围内"
                    f"reason={label_result.get('reason','')}"
                )
                self.total_filtered += 1
                return False, None, f"标签ID {label_id} 不在允许范围内"
        elif label_result and 'skip_reason' in label_result:
            # 模型明确表示不属于任何标签类型
            self.total_filtered += 1
            skip_reason = label_result['skip_reason']
            logger.info(
                f"预标签跳过: unit_id={data.get('unit_id','N/A')}, pr={data.get('pr_number','N/A')}, "
                f"file={data.get('file_path','N/A')} 不属于任何标签类型"
            )
            logger.info(f"模型分类原因: {skip_reason}")
            return False, None, skip_reason
        else:
            # 其他情况（如模型返回格式异常等）
            self.total_filtered += 1
            logger.info(
                f"预标签跳过: unit_id={data.get('unit_id','N/A')}, pr={data.get('pr_number','N/A')}, "
                f"file={data.get('file_path','N/A')} 模型返回异常或无结果"
            )
            return False, None, "模型返回异常或无结果"
    
    def process_data_file(self, input_file: str, output_file: str, log_file: str = None):
        """处理数据文件"""
        logger.info(f"开始处理数据文件: {input_file}")
        
        # 设置日志文件
        if log_file:
            file_handler = logging.FileHandler(log_file, encoding='utf-8')
            file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
            logger.addHandler(file_handler)
        
        processed_data = []
        filtered_data = []
        
        with open(input_file, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                try:
                    data = json.loads(line.strip())
                    
                    # 处理数据
                    is_valid, processed_item, skip_reason = self.process_single_data(data)
                    
                    if is_valid:
                        processed_data.append(processed_item)
                    else:
                        filtered_data.append(data)
                    
                    # 每处理100条数据输出一次进度
                    if line_num % 10 == 0:
                        logger.info(f"已处理 {line_num} 条数据")
                    
                    
                except json.JSONDecodeError as e:
                    logger.error(f"第{line_num}行JSON解析失败: {e}")
                except Exception as e:
                    logger.error(f"处理第{line_num}行时发生错误: {e}")
        
        # 保存处理后的数据
        with open(output_file, 'w', encoding='utf-8') as f:
            for data in processed_data:
                f.write(json.dumps(data, ensure_ascii=False) + '\n')
        
        # 保存过滤的数据（用于调试）
        filtered_file = output_file.replace('.jsonl', '_filtered.jsonl')
        with open(filtered_file, 'w', encoding='utf-8') as f:
            for data in filtered_data:
                f.write(json.dumps(data, ensure_ascii=False) + '\n')
        
        # 输出统计信息
        logger.info(f"数据处理完成！")
        logger.info(f"总处理数据: {self.total_processed}")
        logger.info(f"有效数据: {len(processed_data)}")
        logger.info(f"过滤数据: {len(filtered_data)}")
        logger.info(f"标签分布: {self.label_distribution}")
        
        # 保存统计信息
        stats_file = output_file.replace('.jsonl', '_stats.json')
        stats = {
            "total_processed": self.total_processed,
            "valid_count": len(processed_data),
            "filtered_count": len(filtered_data),
            "label_distribution": self.label_distribution,
            "timestamp": str(time.time())
        }
        
        with open(stats_file, 'w', encoding='utf-8') as f:
            json.dump(stats, f, ensure_ascii=False, indent=2)
        
        return processed_data, filtered_data


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="预标签数据")
    parser.add_argument("--input", "-i", required=True, help="输入文件路径")
    parser.add_argument("--output", "-o", required=True, help="输出文件路径")
    parser.add_argument("--log", "-l", help="日志文件路径")
    parser.add_argument("--api-key", required=True, help="API密钥")
    parser.add_argument("--base-url", required=True, help="API基础URL")
    parser.add_argument("--model", default=Config.get_model_name(), help="模型名称")
    
    args = parser.parse_args()
    
    labeler = PreLabeler(args.api_key, args.base_url, args.model)
    labeler.process_data_file(args.input, args.output, args.log)

if __name__ == "__main__":
    main()
